import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import tushare as ts
import warnings
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score
import talib as ta
from matplotlib.dates import DateFormatter
import matplotlib.ticker as mtick
from matplotlib.colors import LinearSegmentedColormap

# 设置中文字体
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
warnings.filterwarnings('ignore')


class DataProcessor:
    def __init__(self, token):
        ts.set_token(token)
        self.pro = ts.pro_api()
        self.stocks = ['000001.SZ', '600519.SH', '601398.SH', '002594.SZ', '600036.SH']
        self.stock_names = {
            '000001.SZ': '平安银行',
            '600519.SH': '贵州茅台',
            '601398.SH': '工商银行',
            '002594.SZ': '比亚迪',
            '600036.SH': '招商银行'
        }

    def fetch_data(self, start_date, end_date):
        data_dict = {}
        for stock in self.stocks:
            try:
                df = ts.pro_bar(ts_code=stock, adj='qfq',
                                start_date=start_date.replace('-', ''),
                                end_date=end_date.replace('-', ''))

                if df is None or df.empty:
                    print(f"无法获取{self.stock_names.get(stock, stock)}的数据")
                    continue

                df = df.rename(columns={
                    'trade_date': 'Date',
                    'open': 'Open',
                    'high': 'High',
                    'low': 'Low',
                    'close': 'Close',
                    'vol': 'Volume'
                })

                df['Date'] = pd.to_datetime(df['Date'])
                df = df.set_index('Date')
                df = df.sort_index()

                data_dict[stock] = df
                print(f"成功获取{self.stock_names.get(stock, stock)}的数据，共{len(df)}行")

            except Exception as e:
                print(f"获取{self.stock_names.get(stock, stock)}数据时出错: {e}")
        return data_dict

    def preprocess_data(self, data_dict):
        processed_data = {}
        for stock, data in data_dict.items():
            # 基础技术指标
            data['SMA_5'] = ta.SMA(data['Close'], timeperiod=5)
            data['SMA_20'] = ta.SMA(data['Close'], timeperiod=20)
            data['SMA_60'] = ta.SMA(data['Close'], timeperiod=60)
            data['RSI'] = ta.RSI(data['Close'], timeperiod=14)
            data['MACD'], data['MACD_signal'], _ = ta.MACD(data['Close'])
            data['ATR'] = ta.ATR(data['High'], data['Low'], data['Close'], timeperiod=14)
            data['ADX'] = ta.ADX(data['High'], data['Low'], data['Close'], timeperiod=14)

            # K线形态识别
            data['CDLDOJI'] = ta.CDLDOJI(data['Open'], data['High'], data['Low'], data['Close'])
            data['CDLHAMMER'] = ta.CDLHAMMER(data['Open'], data['High'], data['Low'], data['Close'])
            data['CDLHANGINGMAN'] = ta.CDLHANGINGMAN(data['Open'], data['High'], data['Low'], data['Close'])
            data['CDLSHOOTINGSTAR'] = ta.CDLSHOOTINGSTAR(data['Open'], data['High'], data['Low'], data['Close'])

            # 财务指标
            data['Pct_Change'] = data['Close'].pct_change()
            data['Volatility'] = data['Pct_Change'].rolling(window=20).std() * np.sqrt(252)
            data['Momentum'] = data['Close'].pct_change(5)
            data['Volume_Change'] = data['Volume'].pct_change()

            # 计算未来收益率和标签
            data['Next_Return'] = data['Close'].pct_change(1).shift(-1)
            data['Next_5day_Return'] = data['Close'].pct_change(5).shift(-5)
            data['Label'] = np.where(data['Next_Return'] > 0, 1, 0)
            data['Label_5day'] = np.where(data['Next_5day_Return'] > 0, 1, 0)

            # 移除缺失值
            data = data.dropna()
            processed_data[stock] = data
        return processed_data

    def visualize_indicators(self, data, stock_code):
        stock_name = self.stock_names.get(stock_code, stock_code)
        fig, axes = plt.subplots(3, 1, figsize=(14, 12), sharex=True)
        fig.suptitle(f'{stock_name} 价格与技术指标分析', fontsize=16, y=0.99)

        # 绘制价格和均线
        axes[0].plot(data.index, data['Close'], label='收盘价', color='#007BFF', linewidth=2)
        axes[0].plot(data.index, data['SMA_5'], label='5日均线', color='#FF7D00', linestyle='--')
        axes[0].plot(data.index, data['SMA_20'], label='20日均线', color='#22AA22', linestyle='--')
        axes[0].plot(data.index, data['SMA_60'], label='60日均线', color='#AA2222', linestyle='--')
        axes[0].set_title('价格与均线', fontsize=14)
        axes[0].legend(loc='upper left')
        axes[0].grid(True, linestyle='--', alpha=0.7)
        axes[0].yaxis.set_major_formatter(mtick.StrMethodFormatter('{x:.2f}'))

        # 绘制RSI指标
        axes[1].plot(data.index, data['RSI'], label='RSI', color='#007BFF', linewidth=2)
        axes[1].axhline(y=70, color='#AA2222', linestyle='-', alpha=0.7, label='超买线')
        axes[1].axhline(y=30, color='#22AA22', linestyle='-', alpha=0.7, label='超卖线')
        axes[1].fill_between(data.index, data['RSI'], 70, where=(data['RSI'] >= 70),
                             color='#FFCCCC', alpha=0.5)
        axes[1].fill_between(data.index, data['RSI'], 30, where=(data['RSI'] <= 30),
                             color='#CCFFCC', alpha=0.5)
        axes[1].set_title('RSI指标', fontsize=14)
        axes[1].set_ylim(0, 100)
        axes[1].legend(loc='upper left')
        axes[1].grid(True, linestyle='--', alpha=0.7)

        # 绘制MACD指标
        axes[2].plot(data.index, data['MACD'], label='MACD', color='#007BFF', linewidth=2)
        axes[2].plot(data.index, data['MACD_signal'], label='Signal', color='#FF7D00', linewidth=2)
        axes[2].bar(data.index, data['MACD'] - data['MACD_signal'], label='Histogram',
                    color=np.where(data['MACD'] - data['MACD_signal'] >= 0, '#22AA22', '#AA2222'),
                    alpha=0.7)
        axes[2].set_title('MACD指标', fontsize=14)
        axes[2].legend(loc='upper left')
        axes[2].grid(True, linestyle='--', alpha=0.7)

        # 设置x轴日期格式
        date_format = DateFormatter('%Y-%m-%d')
        axes[2].xaxis.set_major_formatter(date_format)
        plt.xticks(rotation=45)

        plt.tight_layout(rect=[0, 0, 1, 0.96])  # 为suptitle留出空间
        plt.show()


class Strategy:
    def __init__(self, data, stock_names):
        self.data = data
        self.stock_names = stock_names

    def price_volume_strategy(self):
        signals = {}
        for stock, df in self.data.items():
            df = df.copy()
            df['Signal'] = 0

            buy_condition = (df['RSI'] < 35)
            sell_condition = (df['RSI'] > 65)

            df.loc[buy_condition, 'Signal'] = 1
            df.loc[sell_condition, 'Signal'] = -1

            # 打印信号统计
            print(
                f"\n{self.stock_names.get(stock, stock)} 信号统计 - 买入: {len(df[df['Signal'] == 1])}, 卖出: {len(df[df['Signal'] == -1])}")
            signals[stock] = df
        return signals

    def financial_strategy(self):
        signals = {}
        for stock, df in self.data.items():
            df = df.copy()
            df['Signal'] = 0

            # 仅在趋势市交易（ADX>25）
            trend_condition = df['ADX'] > 25

            buy_condition = (df['Volatility'] < df['Volatility'].rolling(10).mean()) & \
                            (df['Momentum'] > 0) & \
                            (df['RSI'] < 55) & \
                            trend_condition

            sell_condition = (df['Volatility'] > df['Volatility'].rolling(10).mean()) & \
                             (df['Momentum'] < 0) & \
                             (df['RSI'] > 45) & \
                             trend_condition

            df.loc[buy_condition, 'Signal'] = 1
            df.loc[sell_condition, 'Signal'] = -1
            signals[stock] = df
        return signals

    def ml_strategy(self):
        signals = {}
        feature_importance = {}

        for stock, df in self.data.items():
            df = df.copy()

            # 特征选择
            features = ['SMA_5', 'SMA_20', 'SMA_60', 'RSI', 'MACD', 'MACD_signal',
                        'ATR', 'ADX', 'Volatility', 'Momentum', 'Volume_Change',
                        'CDLDOJI', 'CDLHAMMER', 'CDLHANGINGMAN', 'CDLSHOOTINGSTAR']

            X = df[features]
            y = df['Label_5day']  # 使用5日预测标签

            # 训练测试分割
            train_size = int(len(X) * 0.7)
            X_train, X_test = X.iloc[:train_size], X.iloc[train_size:]
            y_train, y_test = y.iloc[:train_size], y.iloc[train_size:]

            # 模型训练
            model = RandomForestClassifier(n_estimators=100, random_state=42,
                                           class_weight='balanced')
            model.fit(X_train, y_train)

            # 特征重要性
            feature_importance[stock] = pd.Series(model.feature_importances_, index=features)

            # 预测
            df.loc[X_test.index, 'Prediction'] = model.predict(X_test)
            df.loc[X_test.index, 'Probability'] = model.predict_proba(X_test)[:, 1]

            # 生成信号
            df['Signal'] = 0
            df.loc[(df['Prediction'] == 1) & (df['Probability'] > 0.6), 'Signal'] = 1
            df.loc[(df['Prediction'] == 0) & (df['Probability'] > 0.6), 'Signal'] = -1

            # 评估模型
            accuracy = accuracy_score(y_test, df.loc[X_test.index, 'Prediction'])
            print(f"{self.stock_names.get(stock, stock)} 机器学习模型准确率: {accuracy:.2f}")

            signals[stock] = df

        # 可视化特征重要性
        self._visualize_feature_importance(feature_importance)

        return signals

    def _visualize_feature_importance(self, feature_importance):
        n_stocks = len(feature_importance)
        fig, axes = plt.subplots(1, n_stocks, figsize=(5 * n_stocks, 5))

        if n_stocks == 1:
            axes = [axes]

        # 创建自定义颜色映射
        colors = plt.cm.viridis(np.linspace(0.3, 0.9, len(feature_importance)))

        for i, (stock, importance) in enumerate(feature_importance.items()):
            importance = importance.sort_values(ascending=True)
            axes[i].barh(importance.index, importance.values, color=colors[i])
            axes[i].set_title(f"{self.stock_names.get(stock, stock)} 特征重要性", fontsize=12)
            axes[i].grid(axis='x', linestyle='--', alpha=0.7)
            axes[i].set_xlabel('重要性')

        plt.tight_layout()
        plt.show()


class Backtest:
    def __init__(self, signals_dict, stock_names, initial_capital=100000):
        self.signals_dict = signals_dict
        self.initial_capital = initial_capital
        self.stock_names = stock_names

    def run_backtest(self, stop_loss_pct=0.05, take_profit_pct=0.1, trailing_stop_pct=0.03):
        results = {}
        for stock, df in self.signals_dict.items():
            df = df.copy()
            portfolio = pd.DataFrame(index=df.index)
            portfolio['positions'] = 0
            portfolio['cash'] = self.initial_capital
            portfolio['total'] = self.initial_capital

            position = 0
            entry_price = 0
            highest_price_since_entry = 0

            trades = []

            for i in range(1, len(df)):
                signal = df['Signal'].iloc[i]
                close = df['Close'].iloc[i]
                atr = df['ATR'].iloc[i]  # 添加ATR数据

                # 动态止损：2倍ATR
                stop_loss_pct = 2 * atr / close

                if position > 0:
                    # 动态止损
                    if close <= entry_price * (1 - stop_loss_pct):
                        signal = -1
                        trades.append(('动态止损', df.index[i], close, entry_price, close / entry_price - 1))

                # 更新自买入以来的最高价（用于追踪止损）
                if position > 0 and close > highest_price_since_entry:
                    highest_price_since_entry = close

                # 止损止盈逻辑
                if position > 0:
                    # 固定止损
                    if close <= entry_price * (1 - stop_loss_pct):
                        signal = -1  # 止损卖出
                        trades.append(('止损卖出', df.index[i], close, entry_price, close / entry_price - 1))

                    # 止盈
                    elif close >= entry_price * (1 + take_profit_pct):
                        signal = -1  # 止盈卖出
                        trades.append(('止盈卖出', df.index[i], close, entry_price, close / entry_price - 1))

                    # 追踪止损
                    elif highest_price_since_entry > 0 and close <= highest_price_since_entry * (1 - trailing_stop_pct):
                        signal = -1  # 追踪止损卖出
                        trades.append(('追踪止损卖出', df.index[i], close, entry_price, close / entry_price - 1))

                # 交易信号执行
                if signal == 1 and position == 0:  # 买入
                    shares_to_buy = portfolio['cash'].iloc[i - 1] // close
                    portfolio.loc[portfolio.index[i], 'positions'] = shares_to_buy
                    portfolio.loc[portfolio.index[i], 'cash'] = portfolio['cash'].iloc[i - 1] - shares_to_buy * close
                    position = shares_to_buy
                    entry_price = close
                    highest_price_since_entry = close
                    trades.append(('买入', df.index[i], close, None, None))

                elif signal == -1 and position > 0:  # 卖出
                    portfolio.loc[portfolio.index[i], 'positions'] = 0
                    portfolio.loc[portfolio.index[i], 'cash'] = portfolio['cash'].iloc[i - 1] + position * close
                    position = 0
                    entry_price = 0
                    highest_price_since_entry = 0
                    trades.append(('信号卖出', df.index[i], close, entry_price, close / entry_price - 1))

                else:  # 保持当前状态
                    portfolio.loc[portfolio.index[i], 'positions'] = portfolio['positions'].iloc[i - 1]
                    portfolio.loc[portfolio.index[i], 'cash'] = portfolio['cash'].iloc[i - 1]

                # 计算总资产
                portfolio.loc[portfolio.index[i], 'total'] = portfolio['cash'].iloc[i] + portfolio['positions'].iloc[
                    i] * close

            # 计算收益率
            df['Portfolio'] = portfolio['total']
            df['Daily_Return'] = df['Portfolio'].pct_change()
            df['Cumulative_Return'] = (1 + df['Daily_Return']).cumprod()

            # 计算评估指标
            total_return = (df['Portfolio'].iloc[-1] / self.initial_capital) - 1
            annual_return = (1 + total_return) ** (252 / len(df)) - 1

            # 修复：检查标准差是否为零，避免除零错误
            daily_return_std = df['Daily_Return'].std()
            if daily_return_std == 0:
                sharpe_ratio = 0
            else:
                sharpe_ratio = df['Daily_Return'].mean() / daily_return_std * np.sqrt(252)

            # 最大回撤
            df['Cumulative_Max'] = df['Portfolio'].cummax()
            df['Drawdown'] = (df['Portfolio'] - df['Cumulative_Max']) / df['Cumulative_Max']
            max_drawdown = df['Drawdown'].min()

            # 胜率
            if trades:
                profit_trades = [t for t in trades if t[0] in ['止盈卖出', '追踪止损卖出', '信号卖出'] and t[4] > 0]
                win_rate = len(profit_trades) / len(
                    [t for t in trades if t[0] in ['止盈卖出', '追踪止损卖出', '信号卖出']])
            else:
                win_rate = 0

            results[stock] = {
                'total_return': total_return,
                'annual_return': annual_return,
                'sharpe_ratio': sharpe_ratio,
                'max_drawdown': max_drawdown,
                'win_rate': win_rate,
                'portfolio': portfolio,
                'data': df,
                'trades': trades
            }

            print(f"\n{self.stock_names.get(stock, stock)} 回测结果:")
            print(f"总收益率: {total_return:.2%}")
            print(f"年化收益率: {annual_return:.2%}")
            print(f"夏普比率: {sharpe_ratio:.2f}")
            print(f"最大回撤: {max_drawdown:.2%}")
            print(f"胜率: {win_rate:.2%}")
            print(f"交易次数: {len(trades)}")

        return results

    def plot_results(self, results, strategy_name):
        # 创建自定义颜色映射
        colors = plt.cm.tab10.colors

        plt.figure(figsize=(16, 12))
        plt.suptitle(f'{strategy_name} 回测结果分析', fontsize=16, y=0.99)

        # 1. 绘制各股票的策略收益对比
        plt.subplot(2, 2, 1)
        for i, (stock, result) in enumerate(results.items()):
            plt.plot(result['data']['Cumulative_Return'], label=f"{self.stock_names.get(stock, stock)}",
                     color=colors[i], linewidth=1.5)

        plt.plot(
            results[list(results.keys())[0]]['data']['Close'] / results[list(results.keys())[0]]['data']['Close'].iloc[
                0],
            'k--', label='基准收益', linewidth=1.5)
        plt.title('不同股票的累计收益率对比', fontsize=14)
        plt.xlabel('日期', fontsize=12)
        plt.ylabel('累计收益率', fontsize=12)
        plt.legend(loc='upper left')
        plt.grid(True, linestyle='--', alpha=0.7)
        plt.gca().yaxis.set_major_formatter(mtick.PercentFormatter(1.0))

        # 2. 绘制策略与基准的收益对比
        plt.subplot(2, 2, 2)
        strategy_returns = pd.DataFrame()
        benchmark_returns = pd.DataFrame()

        for stock, result in results.items():
            strategy_returns[stock] = result['data']['Daily_Return']
            benchmark_returns[stock] = result['data']['Close'].pct_change()

        strategy_cumulative = (1 + strategy_returns.mean(axis=1)).cumprod()
        benchmark_cumulative = (1 + benchmark_returns.mean(axis=1)).cumprod()

        plt.plot(strategy_cumulative, label='策略平均收益', color='#007BFF', linewidth=2)
        plt.plot(benchmark_cumulative, 'k--', label='基准平均收益', linewidth=2)
        plt.title('策略与基准的平均累计收益率对比', fontsize=14)
        plt.xlabel('日期', fontsize=12)
        plt.ylabel('累计收益率', fontsize=12)
        plt.legend(loc='upper left')
        plt.grid(True, linestyle='--', alpha=0.7)
        plt.gca().yaxis.set_major_formatter(mtick.PercentFormatter(1.0))

        # 3. 绘制回撤曲线
        plt.subplot(2, 2, 3)
        for i, (stock, result) in enumerate(results.items()):
            plt.plot(result['data']['Drawdown'], label=f"{self.stock_names.get(stock, stock)}",
                     color=colors[i], linewidth=1.5)

        plt.title('不同股票的回撤曲线', fontsize=14)
        plt.xlabel('日期', fontsize=12)
        plt.ylabel('回撤幅度', fontsize=12)
        plt.legend(loc='lower left')
        plt.grid(True, linestyle='--', alpha=0.7)
        plt.gca().yaxis.set_major_formatter(mtick.PercentFormatter(1.0))

        # 4. 绘制绩效指标雷达图
        ax = plt.subplot(2, 2, 4, polar=True)
        metrics = ['总收益率', '年化收益率', '夏普比率', '最大回撤', '胜率']
        angles = np.linspace(0, 2 * np.pi, len(metrics), endpoint=False).tolist()

        # 闭合雷达图
        angles += angles[:1]

        for i, (stock, result) in enumerate(results.items()):
            values = [
                result['total_return'],
                result['annual_return'],
                result['sharpe_ratio'] if result['sharpe_ratio'] < 5 else 5,  # 限制夏普比率范围
                -result['max_drawdown'],  # 转为正值便于比较
                result['win_rate']
            ]

            # 闭合雷达图
            values += values[:1]

            ax.plot(angles, values, linewidth=2, label=f"{self.stock_names.get(stock, stock)}", color=colors[i])
            ax.fill(angles, values, alpha=0.1, color=colors[i])

        ax.set_thetagrids(np.degrees(angles[:-1]), metrics, fontsize=12)
        ax.set_ylim(0, 1)
        ax.set_title(f'{strategy_name} - 绩效指标雷达图', fontsize=14, y=1.05)
        ax.legend(loc='upper right', bbox_to_anchor=(1.3, 1.1))
        ax.grid(True, linestyle='--', alpha=0.7)

        plt.tight_layout(rect=[0, 0, 1, 0.96])  # 为suptitle留出空间
        plt.show()

    def visualize_trades(self, results, stock_code):
        stock_name = self.stock_names.get(stock_code, stock_code)
        result = results[stock_code]
        df = result['data']

        fig, axes = plt.subplots(2, 1, figsize=(16, 12), sharex=True)
        fig.suptitle(f'{stock_name} 交易信号与资产曲线', fontsize=16, y=0.99)

        # 绘制价格和交易信号
        axes[0].plot(df.index, df['Close'], label='收盘价', color='#007BFF', linewidth=2)
        axes[0].plot(df.index, df['SMA_20'], label='20日均线', color='#FF7D00', linestyle='--', linewidth=1.5)

        # 标记买入卖出点
        buy_points = df[df['Signal'] == 1]
        sell_points = df[df['Signal'] == -1]

        axes[0].scatter(buy_points.index, buy_points['Close'], marker='^', color='#22AA22', s=100,
                        label='买入信号', alpha=0.8)
        axes[0].scatter(sell_points.index, sell_points['Close'], marker='v', color='#AA2222', s=100,
                        label='卖出信号', alpha=0.8)

        # 标记实际交易
        trades = result['trades']
        buy_trades = [t for t in trades if t[0] == '买入']
        sell_trades = [t for t in trades if t[0] in ['止盈卖出', '追踪止损卖出', '信号卖出']]

        axes[0].scatter([t[1] for t in buy_trades], [t[2] for t in buy_trades], marker='^', color='#00FF00', s=60,
                        label='实际买入', alpha=0.8)
        axes[0].scatter([t[1] for t in sell_trades], [t[2] for t in sell_trades], marker='v', color='#FF0000', s=60,
                        label='实际卖出', alpha=0.8)

        # 绘制交易连接线
        for trade in trades:
            if trade[0] == '买入':
                buy_date = trade[1]
                buy_price = trade[2]
                # 查找对应的卖出交易
                for sell_trade in sell_trades:
                    if sell_trade[1] > buy_date:
                        sell_date = sell_trade[1]
                        sell_price = sell_trade[2]
                        profit_pct = sell_trade[4]
                        color = '#00FF00' if profit_pct > 0 else '#FF0000'
                        axes[0].plot([buy_date, sell_date], [buy_price, sell_price], '--', color=color, alpha=0.3)
                        break

        axes[0].set_title('价格与交易信号', fontsize=14)
        axes[0].set_ylabel('价格', fontsize=12)
        axes[0].legend(loc='upper left')
        axes[0].grid(True, linestyle='--', alpha=0.7)

        # 绘制资产曲线
        axes[1].plot(df.index, df['Portfolio'], label='总资产', color='#007BFF', linewidth=2)
        axes[1].axhline(y=self.initial_capital, color='k', linestyle='--', alpha=0.7, label='初始资金')

        # 标记交易点的资产变化
        portfolio_values = [df['Portfolio'].loc[trade[1]] for trade in trades if trade[0] in ['买入', '信号卖出']]
        trade_dates = [trade[1] for trade in trades if trade[0] in ['买入', '信号卖出']]

        axes[1].scatter(trade_dates, portfolio_values, color='#FF7D00', s=50, alpha=0.7)

        # 添加资产变化百分比标注
        for i, (date, value) in enumerate(zip(trade_dates, portfolio_values)):
            if i > 0:
                pct_change = (value / portfolio_values[i - 1] - 1) * 100
                axes[1].annotate(f"{pct_change:.2f}%", (date, value),
                                 textcoords="offset points",
                                 xytext=(0, 10),
                                 ha='center',
                                 fontsize=9,
                                 color='#FF7D00')

        axes[1].set_title('资产曲线', fontsize=14)
        axes[1].set_xlabel('日期', fontsize=12)
        axes[1].set_ylabel('资产价值', fontsize=12)
        axes[1].legend(loc='upper left')
        axes[1].grid(True, linestyle='--', alpha=0.7)

        # 设置x轴日期格式
        date_format = DateFormatter('%Y-%m-%d')
        axes[1].xaxis.set_major_formatter(date_format)
        plt.xticks(rotation=45)

        plt.tight_layout(rect=[0, 0, 1, 0.96])  # 为suptitle留出空间
        plt.show()

def main():
    token = '1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c'

    # 数据处理
    data_processor = DataProcessor(token)
    raw_data = data_processor.fetch_data(start_date='2019-01-01', end_date='2025-01-01')

    if not raw_data:
        print("没有成功获取任何股票数据，程序退出")
        return

    processed_data = data_processor.preprocess_data(raw_data)

    # 可视化部分股票的指标
    sample_stocks = list(processed_data.keys())[:5]
    for stock in sample_stocks:
        data_processor.visualize_indicators(processed_data[stock], stock)

    # 策略实现
    strategy = Strategy(processed_data, data_processor.stock_names)
    price_volume_signals = strategy.price_volume_strategy()
    financial_signals = strategy.financial_strategy()
    ml_signals = strategy.ml_strategy()

    # 回测
    print("\n=== 量价策略回测 ===")
    price_volume_backtest = Backtest(price_volume_signals, data_processor.stock_names)
    price_volume_results = price_volume_backtest.run_backtest(stop_loss_pct=0.05, take_profit_pct=0.15,
                                                              trailing_stop_pct=0.03)
    price_volume_backtest.plot_results(price_volume_results, '量价策略')

    # 可视化交易细节
    for stock in sample_stocks:
        price_volume_backtest.visualize_trades(price_volume_results, stock)

    print("\n=== 财务指标策略回测 ===")
    financial_backtest = Backtest(financial_signals, data_processor.stock_names)
    financial_results = financial_backtest.run_backtest(stop_loss_pct=0.07, take_profit_pct=0.12,
                                                        trailing_stop_pct=0.04)
    financial_backtest.plot_results(financial_results, '财务指标策略')

    print("\n=== 机器学习策略回测 ===")
    ml_backtest = Backtest(ml_signals, data_processor.stock_names)
    ml_results = ml_backtest.run_backtest(stop_loss_pct=0.08, take_profit_pct=0.15, trailing_stop_pct=0.05)
    ml_backtest.plot_results(ml_results, '机器学习策略')

    # 策略对比
    plt.figure(figsize=(15, 10))
    plt.suptitle('不同策略收益对比', fontsize=16, y=0.99)

    for i, stock in enumerate(processed_data.keys(), 1):
        plt.subplot(2, 3, i)
        plt.plot(price_volume_results[stock]['data']['Cumulative_Return'], label='量价策略',
                 color='#007BFF', linewidth=1.5)
        plt.plot(financial_results[stock]['data']['Cumulative_Return'], label='财务指标策略',
                 color='#FF7D00', linewidth=1.5)
        plt.plot(ml_results[stock]['data']['Cumulative_Return'], label='机器学习策略',
                 color='#22AA22', linewidth=1.5)
        plt.plot(ml_results[stock]['data']['Close'] / ml_results[stock]['data']['Close'].iloc[0], 'k--', label='基准',
                 linewidth=1.5)
        plt.title(f'{data_processor.stock_names.get(stock, stock)}', fontsize=14)
        plt.legend(loc='upper left')
        plt.grid(True, linestyle='--', alpha=0.7)
        plt.gca().yaxis.set_major_formatter(mtick.PercentFormatter(1.0))

    plt.tight_layout(rect=[0, 0, 1, 0.96])  # 为suptitle留出空间
    plt.show()

    # 综合评估
    strategies = ['量价策略', '财务指标策略', '机器学习策略']
    all_results = [price_volume_results, financial_results, ml_results]

    eval_df = pd.DataFrame(
        columns=['股票', '策略', '总收益率', '年化收益率', '夏普比率', '最大回撤', '胜率', '交易次数'])
    for i, strategy_results in enumerate(all_results):
        for stock, result in strategy_results.items():
            eval_df = pd.concat([eval_df, pd.DataFrame({
                '股票': [data_processor.stock_names.get(stock, stock)],
                '策略': [strategies[i]],
                '总收益率': [f"{result['total_return']:.2%}"],
                '年化收益率': [f"{result['annual_return']:.2%}"],
                '夏普比率': [f"{result['sharpe_ratio']:.2f}"],
                '最大回撤': [f"{result['max_drawdown']:.2%}"],
                '胜率': [f"{result['win_rate']:.2%}"],
                '交易次数': [len(result['trades'])]
            })], ignore_index=True)

    print("\n=== 策略综合评估 ===")
    print(eval_df)

    # 可视化策略平均绩效对比
    strategy_performance = pd.DataFrame(
        columns=['策略', '平均总收益率', '平均年化收益率', '平均夏普比率', '平均最大回撤', '平均胜率'])

    for i, strategy_name in enumerate(strategies):
        results = all_results[i]
        avg_total_return = np.mean([result['total_return'] for result in results.values()])
        avg_annual_return = np.mean([result['annual_return'] for result in results.values()])
        avg_sharpe = np.mean([result['sharpe_ratio'] for result in results.values()])
        avg_drawdown = np.mean([result['max_drawdown'] for result in results.values()])
        avg_win_rate = np.mean([result['win_rate'] for result in results.values()])

        strategy_performance = pd.concat([strategy_performance, pd.DataFrame({
            '策略': [strategy_name],
            '平均总收益率': [f"{avg_total_return:.2%}"],
            '平均年化收益率': [f"{avg_annual_return:.2%}"],
            '平均夏普比率': [f"{avg_sharpe:.2f}"],
            '平均最大回撤': [f"{avg_drawdown:.2%}"],
            '平均胜率': [f"{avg_win_rate:.2%}"]
        })], ignore_index=True)

    print("\n=== 策略平均绩效对比 ===")
    print(strategy_performance)

    # 可视化策略绩效对比雷达图
    plt.figure(figsize=(10, 8))
    metrics = ['总收益率', '年化收益率', '夏普比率', '最大回撤', '胜率']
    angles = np.linspace(0, 2 * np.pi, len(metrics), endpoint=False).tolist()

    # 创建自定义颜色映射
    strategy_colors = ['#007BFF', '#FF7D00', '#22AA22']

    for i, strategy_name in enumerate(strategies):
        results = all_results[i]
        values = [
            np.mean([result['total_return'] for result in results.values()]),
            np.mean([result['annual_return'] for result in results.values()]),
            np.mean([result['sharpe_ratio'] for result in results.values()]) if np.mean(
                [result['sharpe_ratio'] for result in results.values()]) < 5 else 5,
            -np.mean([result['max_drawdown'] for result in results.values()]),
            np.mean([result['win_rate'] for result in results.values()])
        ]

        # 闭合雷达图
        values_closed = values + values[:1]
        angles_closed = angles + angles[:1]

        plt.polar(angles_closed, values_closed, linewidth=2, label=strategy_name, color=strategy_colors[i])
        plt.fill(angles_closed, values_closed, alpha=0.1, color=strategy_colors[i])

    plt.title('策略绩效对比雷达图', fontsize=16)
    plt.thetagrids(np.degrees(angles), metrics, fontsize=12)
    plt.ylim(0, 1)
    plt.legend(loc='upper right', bbox_to_anchor=(1.3, 1.1))
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    main()